from typing import Dict, Any, TypedDict, Annotated
from langgraph.graph import StateGraph, END
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.tools import tool
import requests
import json
import datetime

class ToolCallingState(TypedDict):
    messages: Annotated[list, "对话消息列表"]
    result: Annotated[str, "处理结果"]
    tool_results: Annotated[Dict[str, Any], "工具调用结果"]
    current_tool: Annotated[str, "当前使用的工具"]

class ToolCallingDemo:
    """工具调用演示"""
    
    def __init__(self):
        self.graph = self._build_graph()
    
    def _build_graph(self):
        """构建工具调用图"""
        
        workflow = StateGraph(ToolCallingState)
        
        # 添加节点
        workflow.add_node("analyze_request", self._analyze_request)
        workflow.add_node("call_weather_tool", self._call_weather_tool)
        workflow.add_node("call_calculator_tool", self._call_calculator_tool)
        workflow.add_node("call_translator_tool", self._call_translator_tool)
        workflow.add_node("combine_tool_results", self._combine_tool_results)
        workflow.add_node("format_response", self._format_response)
        
        # 设置入口点
        workflow.set_entry_point("analyze_request")
        
        # 设置边
        workflow.add_edge("analyze_request", "call_weather_tool")
        workflow.add_edge("call_weather_tool", "call_calculator_tool")
        workflow.add_edge("call_calculator_tool", "call_translator_tool")
        workflow.add_edge("call_translator_tool", "combine_tool_results")
        workflow.add_edge("combine_tool_results", "format_response")
        workflow.add_edge("format_response", END)
        
        return workflow.compile()
    
    def _analyze_request(self, state: ToolCallingState) -> ToolCallingState:
        """分析请求"""
        messages = state.get("messages", [])
        if not messages:
            messages = [HumanMessage(content="查询天气信息")]
        
        content = messages[-1].content
        
        return {
            "messages": messages,
            "result": f"分析请求: {content}",
            "tool_results": {},
            "current_tool": "analyzer"
        }
    
    def _call_weather_tool(self, state: ToolCallingState) -> ToolCallingState:
        """调用天气工具"""
        try:
            # 模拟天气API调用
            weather_data = {
                "location": "北京",
                "temperature": "25°C",
                "condition": "晴天",
                "humidity": "60%",
                "timestamp": datetime.datetime.now().isoformat()
            }
            
            tool_results = state.get("tool_results", {})
            tool_results["weather"] = weather_data
            
            return {
                "messages": state["messages"],
                "result": state["result"],
                "tool_results": tool_results,
                "current_tool": "weather"
            }
        except Exception as e:
            tool_results = state.get("tool_results", {})
            tool_results["weather"] = {"error": str(e)}
            
            return {
                "messages": state["messages"],
                "result": state["result"],
                "tool_results": tool_results,
                "current_tool": "weather"
            }
    
    def _call_calculator_tool(self, state: ToolCallingState) -> ToolCallingState:
        """调用计算器工具"""
        try:
            # 模拟计算器工具
            calculation_data = {
                "expression": "2 + 3 * 4",
                "result": 14,
                "steps": ["2 + 3 * 4", "2 + 12", "14"],
                "timestamp": datetime.datetime.now().isoformat()
            }
            
            tool_results = state.get("tool_results", {})
            tool_results["calculator"] = calculation_data
            
            return {
                "messages": state["messages"],
                "result": state["result"],
                "tool_results": tool_results,
                "current_tool": "calculator"
            }
        except Exception as e:
            tool_results = state.get("tool_results", {})
            tool_results["calculator"] = {"error": str(e)}
            
            return {
                "messages": state["messages"],
                "result": state["result"],
                "tool_results": tool_results,
                "current_tool": "calculator"
            }
    
    def _call_translator_tool(self, state: ToolCallingState) -> ToolCallingState:
        """调用翻译工具"""
        try:
            # 模拟翻译工具
            translation_data = {
                "source_text": "Hello, how are you?",
                "target_language": "中文",
                "translated_text": "你好，你好吗？",
                "confidence": 0.95,
                "timestamp": datetime.datetime.now().isoformat()
            }
            
            tool_results = state.get("tool_results", {})
            tool_results["translator"] = translation_data
            
            return {
                "messages": state["messages"],
                "result": state["result"],
                "tool_results": tool_results,
                "current_tool": "translator"
            }
        except Exception as e:
            tool_results = state.get("tool_results", {})
            tool_results["translator"] = {"error": str(e)}
            
            return {
                "messages": state["messages"],
                "result": state["result"],
                "tool_results": tool_results,
                "current_tool": "translator"
            }
    
    def _combine_tool_results(self, state: ToolCallingState) -> ToolCallingState:
        """合并工具结果"""
        tool_results = state["tool_results"]
        
        combined_result = "工具调用结果汇总:\n"
        for tool_name, result in tool_results.items():
            if "error" in result:
                combined_result += f"- {tool_name}: 调用失败 - {result['error']}\n"
            else:
                combined_result += f"- {tool_name}: 调用成功\n"
        
        return {
            "messages": state["messages"],
            "result": combined_result,
            "tool_results": tool_results,
            "current_tool": "combiner"
        }
    
    def _format_response(self, state: ToolCallingState) -> ToolCallingState:
        """格式化响应"""
        tool_results = state["tool_results"]
        
        # 创建详细的响应
        response_parts = []
        response_parts.append("=== 工具调用演示结果 ===")
        
        if "weather" in tool_results and "error" not in tool_results["weather"]:
            weather = tool_results["weather"]
            response_parts.append(f"🌤️ 天气信息: {weather['location']} {weather['temperature']} {weather['condition']}")
        
        if "calculator" in tool_results and "error" not in tool_results["calculator"]:
            calc = tool_results["calculator"]
            response_parts.append(f"🧮 计算结果: {calc['expression']} = {calc['result']}")
        
        if "translator" in tool_results and "error" not in tool_results["translator"]:
            trans = tool_results["translator"]
            response_parts.append(f"🌐 翻译结果: {trans['source_text']} → {trans['translated_text']}")
        
        final_response = "\n".join(response_parts)
        
        return {
            "messages": state["messages"] + [AIMessage(content=final_response)],
            "result": final_response,
            "tool_results": tool_results,
            "current_tool": "formatter"
        }
    
    def run(self, input_data: Dict[str, Any], config: Dict[str, Any] = None) -> Dict[str, Any]:
        """运行工具调用图"""
        initial_state = {
            "messages": [HumanMessage(content=input_data.get("message", "查询天气信息"))],
            "result": "",
            "tool_results": {},
            "current_tool": ""
        }
        
        result = self.graph.invoke(initial_state)
        
        return {
            "final_state": result,
            "tool_results": result["tool_results"],
            "tools_used": list(result["tool_results"].keys()),
            "final_result": result["result"]
        }
